import numpy as np
import functools
import random
import time
import multiprocessing as mp
import gc
import psutil
from datetime import datetime
import copy

def print_log(data=None, level='INFO'):
    print(f'{datetime.now().strftime("%Y/%m/%d %H:%M:%S")} {level} {data}')

def cal_relative_diff(real_data, expect_data, diff_thd, type_str='fp16'):
    if 'nan' in str(expect_data) or 'inf' in str(expect_data):
        if type_str.lower() == 'fp16':
            expect_data = 65504
        else:
            expect_data = 3.428e38
    diff = abs(float(real_data) - float(expect_data))
    if diff < diff_thd:
        result = diff
    else:
        result = diff / (float(max(abs(real_data), abs(expect_data))) + 10e-10)
    return result

def cal_relative_diff_np(real_data, expect_data, diff_thd):
    a = np.abs(np.subtract(real_data, expect_data))
    b1 = np.maximum(np.abs(real_data), (np.abs(expect_data)))
    b2 = float((1.0 / (1 << 14)) / diff_thd)
    b = np.add(np.maximum(b1, b2), 10e-10)
    result = np.where(a < diff_thd, a, a / b)
    return result

def cal_nan_inf_diff(real_data, expect_data, output_dtype, diff_abs, diff_thd):
    err_diff = []
    err_idx = []

    if output_dtype == 'fp32':
        inf_value = 3.4028e38
    elif output_dtype == 'bf16':
        inf_value = 3.38e38
    else:
        inf_value = 65504

    real_data_copy = real_data.copy()
    expect_data_copy = expect_data.copy()

    inf_idx = np.where(np.isinf(real_data_copy))[0]
    pos_inf_idx = np.where(real_data_copy[np.isinf(real_data_copy)] > 0)[0]
    neg_inf_idx = np.where(real_data_copy[np.isinf(real_data_copy)] < 0)[0]
    real_data_copy[inf_idx[pos_inf_idx]] = inf_value
    real_data_copy[inf_idx[neg_inf_idx]] = -inf_value

    inf_idx = np.where(np.isinf(expect_data_copy))[0]
    pos_inf_idx = np.where(expect_data_copy[np.isinf(expect_data_copy)] > 0)[0]
    neg_inf_idx = np.where(expect_data_copy[np.isinf(expect_data_copy)] < 0)[0]
    expect_data_copy[inf_idx[pos_inf_idx]] = inf_value
    expect_data_copy[inf_idx[neg_inf_idx]] = -inf_value

    num_idx = np.where(
        ~np.isnan(real_data_copy) & \
        ~np.isinf(real_data_copy) & \
        ~np.isnan(expect_data_copy) & \
        ~np.isinf(expect_data_copy)
    )
    nan_inf_idx = np.setdiff1d(np.arange(len(real_data_copy)), num_idx)

    rdiff = cal_relative_diff_np(
        real_data_copy[num_idx].astype(np.float32),
        expect_data_copy[num_idx].astype(np.float32),
        diff_thd
    )
    num_err_diff = rdiff[rdiff > diff_thd]
    diff_idx_list = num_idx[0]
    num_err_idx = diff_idx_list[np.where(rdiff > diff_thd)]

    _f = lambda x: list(map(str, x[nan_inf_idx].tolist()))
    real_data_str = _f(real_data)
    expect_data_str = _f(expect_data)
    temp_err_idx = np.where(np.array(real_data_str) != np.array(expect_data_str))[0]
    nan_inf_err_idx = nan_inf_idx[temp_err_idx]
    nan_inf_err_diff = diff_abs[nan_inf_err_idx]

    err_idx = num_err_idx.tolist() + nan_inf_err_idx.tolist()
    err_diff = num_err_diff.tolist() + nan_inf_err_diff.tolist()
    
    return np.array(err_diff), np.array(err_idx)

def display_output(real_data, expect_data, start, end, diff_hd, expect_fp32_data=None):
    def format_value(val):
        return 'inf' if 'inf' in str(val) else 'nan' if 'nan' in str(val) else f'{val:.7f}'
    
    def compute_diff(j):
        return cal_relative_diff(expect_data[j], real_data[j], diff_hd)
    
    def print_row(index, exp_fp32, exp_val, real_val, diff_abs, rate_diff):
        if expect_fp32_data is not None:
            print_log(f'{index:08d} \t {format_value(exp_fp32)} \t {format_value(exp_val)}'
                      f'\t {format_value(real_val)} \t {format_value(diff_abs)} \t {rate_diff}')
        else:
            print_log(f'{index:08d} \t {format_value(exp_val)} \t {format_value(real_val)}'
                      f'\t {format_value(diff_abs)} \t {rate_diff}')
            
    separator = '-' * 80
    print_log(separator)
    if expect_fp32_data is not None:
        print_log("Loop \t ExpFP32Out \t ExpFP16Out \t NPUOut \t FpDiff(min) \t RateDiff")
    else:
        print_log("Loop \t ExpectOut \t RealOut \t FpDiff \t RateDiff")
    print_log(separator)

    total_count = int(end - start)
    indices = range(total_count + 1)
    
    if total_count > 20:
        display_indices = list(range(10)) + list(range(total_count - 9, total_count + 1))
    else:
        display_indices = indices

    for i in display_indices:
        if i == 10 and total_count > 20:
            print_log("... \t ... \t ... \t ... \t ...")
        idx = start + i
        expected = expect_data[idx]
        actual = real_data[idx]

        if 'inf' in str(expect_data) or 'nan' in str(expected):
            diff_abs = 'inf' if 'inf' in str(expected) else 'nan'
        else:
            diff_abs = abs(np.float64(expected) - np.float64(actual))
        
        rate_diff = compute_diff(idx)

        exp_fp32 = expect_fp32_data[idx] if expect_fp32_data is not None else None
        print_row(start + i + 1, exp_fp32, expected, actual, diff_abs, rate_diff)

def display_error_output(real_data, expect_data, err_idx, relative_diff):
    def print_row(index, exp_val, real_val, abs_diff, rate_diff):
        print_log(f'{index:08d} \t {exp_val:.7f} \t {real_val:.7f} \t {abs_diff:.7f} \t {rate_diff:.7f}')

    print_log('Error Line--------------------------------------------------------------------------------')
    print_log('Loop \t ExpectOut \t RealOut \t FpDiff \t RateDiff')
    print_log('------------------------------------------------------------------------------------------')

    count = 0
    len_err = len(err_idx)
    for idx in err_idx:
        count += 1
        exp_val = expect_data[idx]
        real_val = real_data[idx]
        abs_diff = abs(np.float64(exp_val) - np.float64(real_val))
        rel_diff = relative_diff[count - 1]

        if count <= 10 or (90 < count <= 100):
            print_row(idx, exp_val, real_val, abs_diff, rel_diff)
        elif count == 11 or (count == 101 and len_err > 100):
            print_log('...     \t ...      \t ...        \t ...   \t  ...')
        elif count > 100:
            break
    
    print_log('Max-RE line:-------------------------------------------------------------------------------')
    max_error = max(relative_diff)
    max_indices = [idx for idx in err_idx if relative_diff[err_idx.index(idx)] == max_error]

    for i, idx in enumerate(max_indices[:3]):
        exp_val = expect_data[idx]
        real_val = real_data[idx]
        abs_diff = abs(np.float64(exp_val) - np.float64(real_val))
        print_row(idx, exp_val, real_val, abs_diff, max_error)
    print_log('-------------------------------------------------------------------------------------------')

def data_compare_np(
    npu_output, 
    cpu_output, 
    diff_thd=0.01, 
    pct_thd=0.05, 
    max_diff_hd=0.1,
    precision_method=0,
    rtol=0.005,
    atol=0.000025,
    output_dtype=None
):
    max_error_idx = 10000000
    real_data = npu_output.flatten()
    data_compe = cpu_output.flatten()
    if real_data.size == 0 and real_data.size == data_compe.size:
        print_log("The npu_output is [], and it is same as bm_output, the result of data_compare is Pass")
        return "Pass", 100.0, 0

    start, end = 0, real_data.size - 1
    if end < start:
        end = start
    max_error = 0
    result = "FAIL"
    if real_data.size != data_compe.size:
        print_log(f"Error, the size of npu output[{real_data.size}] and benchmark[{data_compe.size}] is not equal.")
        return result, 0.0, max_error

    overflows_count = data_compe[np.isinf(data_compe)].size + \
        data_compe[np.isnan(data_compe)].size
    if overflows_count > 0:
        print_log(f'Overflow, size: {overflows_count}, benchmark_output: {data_compe[np.isinf(data_compe)][:10]}, {data_compe[np.isnan(data_compe)][:10]}')
    
    split_count = int(end - start + 1) if end != start else 1
    print_log(f"split_count:{float(split_count)}; max_diff hd:{max_diff_hd};")

    has_nan_inf = np.isfinite(real_data).all() or np.isfinite(data_compe).all()

    if precision_method == 0:
        try:
            diff_abs = np.abs(np.subtract(
                real_data.astype(np.float32), data_compe.astype(np.float32)
            ))
        except MemoryError:
            return result, 0.0, max_error
        if has_nan_inf:
            err_diff, err_idx = cal_nan_inf_diff(real_data, data_compe, output_dtype, diff_abs, max_diff_hd)
        else:
            diff_index = np.where(diff_abs > 0)
            rdiff = cal_relative_diff_np(
                real_data[diff_index].astype(np.float32),
                data_compe[diff_index].astype(np.float32),
                max_diff_hd
            )
            err_diff = rdiff[rdiff > max_diff_hd]
            diff_idx_list = diff_index[0]
            err_idx =  diff_idx_list[np.wher(rdiff > max_diff_hd)]
        
        fullfill_percent = float(split_count - err_diff.size) / float(split_count) * 100.0
        display_output(real_data, data_compe, start, end, max_diff_hd)
        pct_thd = (1 - pct_thd) * 100.0
        result = "PASS" if (fullfill_percent >= pct_thd) else "FAIL"

        if len(err_diff) > 0:
            max_error = max(err_diff[0 : max_error_idx])
            if max_error_idx > max_diff_hd:
                result = "FAIL"
                print("max_error > max_diff_hd")
        print_log('---------------------------------------------------------------------------------------------')
        print_log('DiffThd  \t PctThd   \t PctRlt   \t Result')
        print_log('---------------------------------------------------------------------------------------------')
        print_log(f'{diff_thd:.4f}    \t {pct_thd:.2f}%   \t {fullfill_percent:.6f}%   \t {result}')

        if len(err_diff) > 0:
            print_log(f'Max-RelativeError is: {max_error}. Threshold is: {max_diff_hd}')
        if result == 'FAIL':
            display_error_output(real_data, data_compe, err_idx, err_diff[: max_error_idx])

    if precision_method == 1:
        raise NotImplementedError()

    if precision_method == 3:
        raise NotImplementedError()

    return result, fullfill_percent
    
    
